%% Replication file of Baqaee, Burstein and Koike-Mori (2022) %%%%%%%%%%%%%%
%  Written by Yasutaka Koike-Mori (PhD student in economics at UCLA), 11/11/2022
% 
%  Description: 
%  1. Generate artificial data (Nh-CES(constant,variable)) and these artificial data are stored in "data/DataforMatlab".
%  2. Calculate Money-Metics from Baqaee, Burstein and Koike-Mori (2022)'s algorithm.
%  All matrices are in (I,N,T) order: 
%  I: the number of goods, N is the number of people, and T is the number of time segments.
%
%  functions (at Bottom of the main file)
%   Inputs  I_vec, pvec: exogenous wage and price (I,N,T)
%   Outputs v_vec:indirect util, u_vec: Money metric , B_vec: budget shares
%   [v_vec, u_vec, B_vec] = NhCES(I_vec,pvec,paras) 
 
clear all

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%% Setting: Parameters
% Generating artificial data from non-homothetic preference
for s = 1:3

if s == 1 || s == 2
    paras.flag_sigma = 0; % 0 for constant sigma, 1 for variable sigma
else
    paras.flag_sigma = 1; 
end

if s == 1 || s == 3
    paras.eps_vec = [0.2;1;1.65];% Non-homothetic parametrization from Comin et.al (2021)
else
    paras.eps_vec = [1;1;1];% Homotheric
end

N = 100;
I = 3;% Need to adjust it when the number of goods changed.
T = 40;

% The indirect utility of Nh-CES is calculated numerically from wages and prices.
% To make it easier to find a solution, we solve for indirect utility for Ns people, one for each t.
% Therefore, N must be a multiple of Ns.
Ns = 20;

% Nh-CES parameters
if paras.flag_sigma == 0
    paras.sig = 0.25;% from Comin et.al (2021)
else
    paras.sig = 10;
end
paras.eta = -2;% from Auer et.al. (2021)


paras.T = T;paras.T = T;paras.N = N;paras.I = I;paras.Ns = Ns;

%% Setting: Price and Wage
% Price Schedule
pvec(1,1,:) = reshape(exp(linspace(0,2,T)),[1,1,T]);% (I,1,T)
pvec(2,1,:) = reshape(exp(linspace(0,1.7,T)),[1,1,T]);% (I,1,T)
pvec(3,1,:) = reshape(exp(linspace(0,1.4,T)),[1,1,T]);% (I,1,T)

% Generate artificial wage 
trend = (reshape(exp(linspace(0,log(10),T)),[1,1,T]));
rng(1,'twister')
I_vec_tmp = lognrnd(3.58,0.7,1,N);
I_vec = sort(I_vec_tmp(:,:,1),2).*trend;
paras.uvec_init = I_vec/10; % Initial value used to calculate the indirect utility of Comin.

% Calibration of Omega
% The share parameter is calibrated separately so that the budget shares of each good
% for the median household in the firstperiod are all the same (equal to one third for each good)
medianw = median(I_vec(1,:,1));
omega_tmp = 1./(medianw.^([0.2;1;1.65]*(1-paras.sig)));
omega_init = omega_tmp./sum(omega_tmp);
omega_init2 = omega_init(1:end-1,:);
options = optimoptions('fsolve', 'Algorithm', 'trust-region', 'TolFun', 1e-7,'Display','off');
omega_tmp2   =  fsolve( @(x) CalOmega(x,medianw,paras),omega_init2,options);  
paras.omega = [omega_tmp2;1-sum(omega_tmp2)];
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%% Generating Artificial data %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
[v_vec, u_vec, B_vec]= NhCES(I_vec,pvec,paras);
% Algorithm:
[U_vec] = CalMoneyMetric(I_vec, B_vec, pvec,0);

clearvars -except paras pvec u_vec v_vec I_vec B_vec I N T U_vec s

if s == 1 
    save ../data/DataforMatlab/NhCES_const_artificial
end
if s == 2
    save ../data/DataforMatlab/NhCES_const_homothetic_artificial
end
if s == 3
    save ../data/DataforMatlab/NhCES_var_artificial
end

end

%% 
function [v_vec, u_vec, B_vec]= NhCES(I_vec,pvec,paras) 
%%Nonhomothetic Preference: Given Wfull and P, solve for u 
T = paras.T;
N = paras.N;
I = paras.I;
Ns = paras.Ns;
v_vec = zeros(1,N,T);
uvec_init = paras.uvec_init;
options = optimoptions('fsolve', 'Algorithm', 'trust-region', 'TolFun', 1e-7,'Display','off');

for t = 1:T
    for ns = 1: Ns   
        Index = (N/Ns)*(ns-1)+1:(N/Ns)*(ns);
       [ v_vec(:,Index,t),fval]  =  fsolve( @(x) NHU(x,I_vec(:,Index,t),pvec(:,:,t),paras),uvec_init(:,Index),options);  % pid
    end 
end
%%Obtain EV
[u_vec,B_vec]= CalEV(v_vec,pvec,paras);

function err = NHU(v_vec,I_vec,pvec,paras) 
if paras.flag_sigma ==1
    sig0 = paras.sig;
    paras.sig =sig0+log((v_vec).^paras.eta); 
end
E = sum(paras.omega.*(((v_vec.^paras.eps_vec).*pvec).^(1-paras.sig))).^(1./(1-paras.sig));
err = abs(E-I_vec);
end

function [u_vec,B_vec]= CalEV(v_vec,pvec,paras)
paras.sig =paras.sig;
sig0 = paras.sig;
if paras.flag_sigma ==1
     paras.sig =sig0+log((v_vec).^paras.eta);
end
pvec_b = pvec(:,:,1);
u_vec = sum(paras.omega.*(((v_vec.^paras.eps_vec).*pvec_b).^(1-paras.sig))).^(1./(1-paras.sig));
tmp = paras.omega.*(((v_vec.^paras.eps_vec).*pvec).^(1-paras.sig));
B_vec = tmp./sum(tmp);
end

end


function  err = CalOmega(omega_init,I_vec,paras) 
% function for solve for Omega 
options = optimoptions('fsolve', 'Algorithm', 'trust-region', 'TolFun', 1e-9, 'Display', 'off');
 [ v_vec ]  =  fsolve( @(x) NHU(x,I_vec,omega_init,paras),I_vec/10,options);  
if paras.flag_sigma ==1
    sig0 = paras.sig;
    paras.sig =sig0+log((v_vec).^paras.eta);
end

omega = [omega_init;1-sum(omega_init)]; 
tmp = omega.*(((v_vec.^paras.eps_vec)).^(1-paras.sig));
B_vec = tmp./sum(tmp);

   err = (abs(B_vec(1:end,:)-1/paras.I));
    
function err = NHU(v_vec,I_vec,omega_init,paras) 
if paras.flag_sigma ==1
    sig0 = paras.sig;
    paras.sig =sig0+log((v_vec).^paras.eta);
end
omega = [omega_init;1-sum(omega_init)];
E = sum(omega.*(((v_vec.^paras.eps_vec)).^(1-paras.sig))).^(1./(1-paras.sig));
err = abs(E-I_vec);
end
end